import pandas as pd
from openpyxl import Workbook
from mb_core.client.mongo_client import Mongo
from mb_core.util.common import format_time

from zc_core.util.common import list_of_groups

meta_map = [
    {'col': '_id', 'title': '商品编码'},
    {'col': 'skuCode', 'title': '商品编号'},
    {'col': 'skuName', 'title': '商品名称'},
    {'col': 'salePrice', 'title': '销售价'},
    {'col': 'originPrice', 'title': '市场价'},
    {'col': 'spuId', 'title': '同款编码'},
    {'col': 'supplierId', 'title': '供应商编码'},
    {'col': 'supplierName', 'title': '供应商名称'},
    {'col': 'supplierSkuId', 'title': '供应商商品编码'},
    {'col': 'catalog3Name', 'title': '三级分类'},
]


def export_data(table, file):
    # 初始化数据库链接
    mongo = Mongo()
    # 结构组装
    columns = list()
    headers = list()
    for meta in meta_map:
        columns.append(meta.get('col'))
        headers.append(meta.get('title'))

    # 数据查询
    rows = mongo.list(
        collection=table,
        fields=columns,
        query={},
        sort=[]
    )
    print('总数量：%s' % len(rows))

    sub_list = list_of_groups(rows, 800000)
    print('集合数量：%s' % len(sub_list))

    # pandas库写入Excel
    # with pd.ExcelWriter(file) as writer:
    #     for idx, sub in enumerate(sub_list):
    #         df1 = pd.DataFrame(sub)
    #         df1.to_excel(writer, sheet_name=f'Sheet{idx}', index=False, header=True)
    #         print(f'Sheet{idx} -> {len(sub)}')

    for idx, sub in enumerate(sub_list):
        write = pd.ExcelWriter(file % idx)
        df = pd.DataFrame(sub)
        df.to_excel(write, sheet_name='数据', columns=columns, header=headers, index=False)
        write.save()


    # openpyxl库写入Excel
    # wb = Workbook()
    # ws = wb.active
    # ws.append(headers)
    # for row in rows:
    #     result = list()
    #     for col in columns:
    #         result.append(row.get(col))
    #     ws.append(result)
    # wb.save(file)

    print('over~')


if __name__ == '__main__':
    # batch_no = time_to_batch_no(datetime.now(), delta=-1)
    batch_no = '20211103'
    export_data(f'data_{batch_no}', f'/work/华润守正_{batch_no}_%s.xlsx')
